# *-* coding:utf8 *-*

import sys
from typing import NewType
from PyQt5.QtWidgets import QMainWindow ,QApplication,QMessageBox
from PyQt5.QtGui import QIcon

from select_ui import select_ui
from mov_window import movWindow
from train_a_model import MainWindow as train_ui
#from train_window import train_ui
import sys
from PyQt5.QtWidgets import QMainWindow ,QApplication,QMessageBox

from PyQt5.Qt import QThread,QMutex

import gym
from stable_baselines3 import PPO
from torch._C import set_flush_denormal
import puck_world


## 登录界面
class MainWindow(QMainWindow):
    def __init__(self,parent=None):
        super(MainWindow,self).__init__(parent)
        self.ui = select_ui()
        self.ui.setupUi(self)
        self.setWindowTitle("功能选择")
        # 接入其他两个界面
        self.train_ui =train_ui()
        self.mov_ui = movWindow()

        self.connet_imitation_button()
        self.connet_train_button()
        self.connect_test_button()
        self.exit()
    def show_train_ui(self):
        self.close()
        self.train_ui.show()
    def show_mov_ui(self):
        self.close()
        self.mov_ui.show()


    def check_model(self):
        path = self.ui.lineEdit.text()
        ## 读取文件路径
        if path:
            success_rate = self.test_model(path)
        else:
            path ="/home/maxwene/Project/mov_qt_ui/model"+"/test_model"
            success_rate = self.test_model(path)

        QMessageBox.about(self, "成功率", str(success_rate))
        ## 写一个文件完事
    def test_model(self,model):
        env_test = gym.make("puck-world-v0")
        model_predict = PPO("MlpPolicy", env_test, verbose=1)
        model_predict.load(model)
        obs = env_test.reset()

        print("load model success ")
        success = 1 
        for turn in range(100):
            obs = env_test.reset()
            for i in range(1000):
                action, _states =  model_predict.predict(obs, deterministic=True)
                obs, reward, done, info = env_test.step(action)
                if done:
                    success+=1
                    break
        env_test.close()
        return success/100.0
    def connet_imitation_button(self):
        self.ui.pushButton.clicked.connect(self.show_mov_ui)

    def connet_train_button(self):
        self.ui.pushButton_2.clicked.connect(self.show_mov_ui)

    def connect_test_button(self):
        self.ui.pushButton_4.clicked.connect(self.check_model)

    def exit(self):
        self.ui.pushButton_3.clicked.connect(self.close)


if __name__ =="__main__":
    app = QApplication(sys.argv)
    # app.setWindowIcon(QIcon("./images/cartoon1.ico"))

    B = MainWindow()
    B.show()

    sys.exit(app.exec_())